import numpy as np
from PIL import Image
import cv2
import json
from pathlib import Path
try:
    from scalePcd import scalePcdD
except ImportError:
    from .scalePcd import scalePcdD
import os
import pickle
import open3d as o3d
from termcolor import colored   
# Add torch support for tensor conversion
try:
    import torch
    TORCH_AVAILABLE = True
except ImportError:
    TORCH_AVAILABLE = False

def ensure_dir(file_path):
    file_path=os.path.abspath(file_path)
    directory = os.path.dirname(file_path)
    if not os.path.exists(directory):
        os.makedirs(directory)


def save_img(save_path, img_file,surpress_print=False,print_color="black"):
    # Handle different data types for image saving
    if img_file.dtype == np.float32 or img_file.dtype == np.float64:
        # For float arrays, normalize to 0-255 range and convert to uint8
        if img_file.max() > img_file.min():
            img_normalized = ((img_file - img_file.min()) / 
                            (img_file.max() - img_file.min()) * 255).astype(np.uint8)
        else:
            img_normalized = np.zeros_like(img_file, dtype=np.uint8)
        img = Image.fromarray(img_normalized)
    elif img_file.dtype == np.uint16:
        # For uint16 arrays, normalize to 0-255 range and convert to uint8
        if img_file.max() > img_file.min():
            img_normalized = ((img_file - img_file.min()) / 
                            (img_file.max() - img_file.min()) * 255).astype(np.uint8)
        else:
            img_normalized = np.zeros_like(img_file, dtype=np.uint8)
        img = Image.fromarray(img_normalized)
    elif img_file.dtype == np.uint8:
        # For uint8 arrays, use as is
        img = Image.fromarray(img_file)
    else:
        # For other data types, try to convert to uint8
        try:
            img_normalized = img_file.astype(np.uint8)
            img = Image.fromarray(img_normalized)
        except:
            # If conversion fails, normalize to 0-255 range
            if img_file.max() > img_file.min():
                img_normalized = ((img_file - img_file.min()) / 
                                (img_file.max() - img_file.min()) * 255).astype(np.uint8)
            else:
                img_normalized = np.zeros_like(img_file, dtype=np.uint8)
            img = Image.fromarray(img_normalized)
    if not surpress_print:
        print(colored(f"save_img "+ os.path.abspath(save_path)+" "+ str(img_file.shape) ,print_color))
    ensure_dir(save_path)
    img.save(save_path)


def save_json(save_path, json_file,print_color="black"):
    ensure_dir(save_path)
    
    def convert_numpy_to_list(obj):
        """Recursively convert numpy arrays and torch tensors to lists for JSON serialization"""
        if TORCH_AVAILABLE and hasattr(obj, 'detach'):
            # Handle torch tensors
            if hasattr(obj, 'cpu'):
                obj = obj.detach().cpu().numpy()
            else:
                obj = obj.detach().numpy()
        
        if isinstance(obj, np.ndarray):
            return obj.tolist()
        elif isinstance(obj, (Path, )):
            return str(obj)
        elif isinstance(obj, dict):
            return {key: convert_numpy_to_list(value) for key, value in obj.items()}
        elif isinstance(obj, list):
            return [convert_numpy_to_list(item) for item in obj]
        elif isinstance(obj, tuple):
            return tuple(convert_numpy_to_list(item) for item in obj)
        elif isinstance(obj, (np.integer, np.floating)):
            return obj.item()
        # Handle torch scalar types if available
        elif TORCH_AVAILABLE and hasattr(obj, 'item') and callable(getattr(obj, 'item')):
            try:
                return obj.item()
            except:
                pass
        else:
            return obj

    # Convert numpy arrays to lists
    json_serializable = convert_numpy_to_list(json_file)
    
    def custom_json_dump(obj, fp, **kwargs):
        """Custom JSON dumper that keeps leaf arrays on single lines and trims floats to 3 decimal places"""
        def format_number(num):
            """Format numbers: floats to 3 decimal places (trim trailing zeros), others as is"""
            if isinstance(num, float):
                return ("{:.3f}".format(num)).rstrip('0').rstrip('.')
            return str(num)
        def format_value(value, level=0):
            if isinstance(value, list):
                # Check if this is a leaf array (contains only numbers/strings)
                is_leaf = all(isinstance(item, (int, float, str, bool)) or item is None for item in value)

                if is_leaf:  # Keep all leaf arrays on one line regardless of length
                    return '[' + ', '.join(format_number(item) for item in value) + ']'

                # Multi-line format for non-leaf lists
                indent = '  ' * level
                items = []
                for item in value:
                    items.append(format_value(item, level + 1))
                return '[\n' + ',\n'.join(indent + '  ' + item for item in items) + '\n' + indent + ']'
            elif isinstance(value, dict):
                indent = '  ' * level
                items = []
                for key, val in value.items():
                    if isinstance(val, (list, dict)):
                        items.append(f'"{key}": {format_value(val, level + 1)}')
                    else:
                        items.append(f'"{key}": {format_number(val)}')
                return '{\n' + ',\n'.join(indent + '  ' + item for item in items) + '\n' + indent + '}'
            else:
                return format_number(value)
        
        formatted = format_value(obj)
        fp.write(formatted)
    
    with open(save_path, 'w') as f:
        custom_json_dump(json_serializable, f)


def save_pkl(save_path, dic_file,print_color="black"):
    ensure_dir(save_path)
    with open(save_path, 'wb') as f:
        pickle.dump(dic_file, f)


def save_pcd(save_path, pcd_arr, color=True, surpress_print=False,print_color="black"):
    pcd_arr = scalePcdD(pcd_arr)
    if not surpress_print:
        print(colored(f"save_pcd "+ os.path.abspath(save_path)+" "+ str(pcd_arr.shape) ,print_color))
    ensure_dir(save_path)
    
    # Convert torch tensor to numpy if needed
    if TORCH_AVAILABLE and hasattr(pcd_arr, 'detach'):
        if hasattr(pcd_arr, 'cpu'):
            pcd_arr = pcd_arr.detach().cpu().numpy()
        else:
            pcd_arr = pcd_arr.detach().numpy()
    
    # Check if the array is empty or has wrong shape
    if pcd_arr.size == 0:
        print(colored("Warning: Empty point cloud array, creating empty file",print_color))
        point_cloud = o3d.geometry.PointCloud()
        o3d.io.write_point_cloud(save_path, point_cloud)
        return
    
    # Handle multi-dimensional arrays by flattening to 2D
    if pcd_arr.ndim > 2:
        # For shape like (B, T, N, 6), reshape to (B*T*N, 6)
        original_shape = pcd_arr.shape
        print(colored(f"Reshaping point cloud from {original_shape} to 2D for saving",print_color))
        pcd_arr = pcd_arr.reshape(-1, original_shape[-1])
    
    # Ensure the array is 2D
    if pcd_arr.ndim == 1:
        print(colored("Warning: 1D array received, reshaping to 2D",print_color))
        if pcd_arr.shape[0] % 6 == 0:
            pcd_arr = pcd_arr.reshape(-1, 6)
        elif pcd_arr.shape[0] % 3 == 0:
            pcd_arr = pcd_arr.reshape(-1, 3)
        else:
            print(colored(f"Warning: Cannot reshape 1D array of size {pcd_arr.shape[0]} to point cloud format",print_color))
            return
    
    point_cloud = o3d.geometry.PointCloud()
    point_arr = pcd_arr[:, :3]
    point_cloud.points = o3d.utility.Vector3dVector(point_arr)
    if color and pcd_arr.shape[1] >= 6:
        colors_arr = pcd_arr[:, -3:]
        point_cloud.colors = o3d.utility.Vector3dVector(colors_arr)
    o3d.io.write_point_cloud(save_path, point_cloud)


def save_depth(save_path, depth_array, colormap='jet'):
    """
    Save depth data as a colored image.
    
    Args:
        save_path: Path to save the depth image
        depth_array: Depth data array of shape [H, W] with float values (in millimeters)
        colormap: Colormap to apply to depth data ('jet', 'viridis', 'plasma', 'hot', etc.)
    """
    ensure_dir(save_path)
    
    # Handle invalid depth values (set them to 0)
    depth_clean = np.copy(depth_array)
    depth_clean[np.isnan(depth_clean)] = 0
    depth_clean[np.isinf(depth_clean)] = 0
    
    # Normalize depth values to 0-255 range for visualization
    if depth_clean.max() > depth_clean.min():
        depth_normalized = ((depth_clean - depth_clean.min()) / 
                          (depth_clean.max() - depth_clean.min()) * 255).astype(np.uint8)
    else:
        depth_normalized = np.zeros_like(depth_clean, dtype=np.uint8)
    
    # Apply colormap
    if colormap == 'jet':
        colored_depth = cv2.applyColorMap(depth_normalized, cv2.COLORMAP_JET)
    elif colormap == 'viridis':
        colored_depth = cv2.applyColorMap(depth_normalized, cv2.COLORMAP_VIRIDIS)
    elif colormap == 'plasma':
        colored_depth = cv2.applyColorMap(depth_normalized, cv2.COLORMAP_PLASMA)
    elif colormap == 'hot':
        colored_depth = cv2.applyColorMap(depth_normalized, cv2.COLORMAP_HOT)
    else:
        colored_depth = cv2.applyColorMap(depth_normalized, cv2.COLORMAP_JET)
    
    # Convert BGR to RGB for PIL
    colored_depth_rgb = cv2.cvtColor(colored_depth, cv2.COLOR_BGR2RGB)
    
    # Save as image
    img = Image.fromarray(colored_depth_rgb)
    img.save(save_path)


def save_seg(save_path, segmentation_array):
    """
    Save segmentation data as an image.
    
    Args:
        save_path: Path to save the segmentation image
        segmentation_array: Segmentation data array of shape [H, W, 3] with RGB values (0-255)
    """
    ensure_dir(save_path)
    
    # Ensure the data is in the correct format (uint8)
    if segmentation_array.dtype != np.uint8:
        segmentation_array = segmentation_array.astype(np.uint8)
    
    # Create and save image
    img = Image.fromarray(segmentation_array)
    img.save(save_path)


def save_info(info_name, save_path_without_ext, info, surpress_print=False,print_color="black"):
    ensure_dir(save_path_without_ext)
    full_path = str(save_path_without_ext) + '_' + info_name
    if info_name in ['rgba','rgb','img']:
        full_path += '.png'
        save_img(full_path, info,print_color=print_color)
    elif info_name == 'depth':
        full_path += '.png'
        save_depth(full_path, info,print_color=print_color)
    elif info_name == 'seg':
        full_path += '.png'
        save_seg(full_path, info,print_color=print_color)
    elif info_name == 'pcd':
        full_path += '.pcd'
        save_pcd(full_path, info, surpress_print=surpress_print,print_color=print_color)
    elif info_name in ['config','json']:
        full_path += '.json'
        save_json(full_path, info,print_color=print_color)
    elif info_name == 'pkl':
        full_path += '.pkl'
        save_pkl(full_path, info,print_color=print_color)
    return os.path.abspath(full_path)


def save_ig_pcd(pcd_arr, color=True, title="Point Cloud Visualization", figsize=(10, 8), dpi=100):
    import matplotlib.pyplot as plt
    from mpl_toolkits.mplot3d import Axes3D
    assert TORCH_AVAILABLE and hasattr(pcd_arr, 'detach') or not hasattr(pcd_arr, 'detach'), f"TORCH_AVAILABLE={TORCH_AVAILABLE}, has_detach={hasattr(pcd_arr, 'detach')}"
    pcd_arr = pcd_arr.detach().cpu().numpy() if TORCH_AVAILABLE and hasattr(pcd_arr, 'detach') else pcd_arr
    assert pcd_arr.size > 0, f"pcd_arr.size={pcd_arr.size}"
    if pcd_arr.ndim > 2:
        original_shape = pcd_arr.shape
        pcd_arr = pcd_arr.reshape(-1, original_shape[-1])
    if pcd_arr.ndim == 1:
        pcd_arr = pcd_arr.reshape(-1, 6) if pcd_arr.shape[0] % 6 == 0 else pcd_arr.reshape(-1, 3) if pcd_arr.shape[0] % 3 == 0 else pcd_arr
    points = pcd_arr[:, :3]
    colors = pcd_arr[:, -3:] / 255.0 if color and pcd_arr.shape[1] >= 6 and pcd_arr[:, -3:].max() > 1.0 else pcd_arr[:, -3:] if color and pcd_arr.shape[1] >= 6 else None
    figure = plt.figure(figsize=figsize, dpi=dpi)
    axis = figure.add_subplot(111, projection='3d')
    axis.scatter(points[:, 0], points[:, 1], points[:, 2], c=colors, s=20, alpha=0.8) if colors is not None else axis.scatter(points[:, 0], points[:, 1], points[:, 2], s=20, alpha=0.8, c='blue')
    axis.set_xlabel('X')
    axis.set_ylabel('Y')
    axis.set_zlabel('Z')
    axis.set_title(title)
    axis.set_box_aspect([1, 1, 1])
    plt.tight_layout()
    plt.savefig('debug_pcd.png', dpi=dpi, bbox_inches='tight')
    plt.close()


if __name__ == "__main__":
    save_json("test.json", [[1,2,3],[4,5,6]])
